"""
1. 题目描述：基于Tensorflow框架，构建包含一个卷积层和一个池化层的神经网络，实现数据集cifar10的分类。
卷积层1参数：输出通道数为32，卷积核大小为3*3，padding为same，激活函数为ReLU。
池化层1参数：核大小为（2，2），stride为（2，2）。
"""

# 2. 题目要求：
# ①　导入必要的依赖库。（6分）
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks
import numpy as np
import os

np.random.seed(1)
tf.random.set_seed(1)

VER = 'v1.0'
ALPHA = 0.001
BATCH_SIZE = 64
N_EPOCHS = 4
N_EPOCHS_INTERV = 2
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)


# ②　定义一个load_data的方法，用于读取数据。（7分）
def load_data():
    return keras.datasets.cifar10.load_data()


# ③　定义一个Class CifarData类，用于操作数据。（7分）
class CifarData():

    def __init__(self, x, y, batch):
        self.ds = tf.data.Dataset.from_tensor_slices((x, y)) \
            .shuffle(2000) \
            .batch(batch) \
            .prefetch(tf.data.experimental.AUTOTUNE)

    def get_ds(self):
        return self.ds


# ④　获取所有文件名称和标签。（6分）
# ⑤　获取训练数据和测试数据。（6分）
(x_train, y_train), (x_test, y_test) = load_data()
ds_train = CifarData(x_train, y_train, BATCH_SIZE).get_ds()
ds_test = CifarData(x_test, y_test, BATCH_SIZE).get_ds()

# ⑥　定义输入的数据占位符X和输入的标签占位符Y。（6分）
# ⑦　对图片进行相应的转换处理。（6分）
# ⑧　定义网络模型，网络模型结构为：卷积->池化->全连接层。（7分）
model = keras.Sequential([
    layers.Conv2D(32, (3, 3), (1, 1), 'same'),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D((2, 2), (2, 2), 'same'),

    layers.Conv2D(64, (3, 3), (1, 1), 'same'),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D((2, 2), (2, 2), 'same'),

    layers.Flatten(),
    layers.Dense(1024, activation=activations.relu),
    layers.Dense(256, activation=activations.relu),
    layers.Dense(10, activation=None),
])
model.build(input_shape=(None, 32, 32, 3))
model.summary()

# ⑨　计算损失函数，损失函数采用交叉熵函数。（7分）
# ⑩　计算准确率。（7分）
# 11　进行网络模型优化，优化器采用Adam。（7分）
model.compile(
    loss=losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    metrics=metrics.sparse_categorical_accuracy
)


# 12　设置网络模型初始参数，batch_size = 20, 训练迭代次数为10000，测试迭代次数为100。（7分）
# 13　创建会话，并初始化全局变量。（7分）
# 14　开始训练模型，循环迭代10000次，每500次打印输出一次训练loss；每5000次计算一次测试的准确率。（7分）
class MyCallback(callbacks.Callback):

    def __init__(self, epoch_interv, **kwargs):
        super().__init__(**kwargs)
        self.epoch_interv = epoch_interv

    def on_epoch_end(self, epoch, logs=None):
        if epoch % self.epoch_interv == 0:
            print('-' * 32, f'{epoch + 1}', '-' * 32)
            print(logs)


model.fit(ds_train, epochs=N_EPOCHS,
          validation_data=ds_test,
          callbacks=[
              callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0),
              MyCallback(N_EPOCHS_INTERV),
          ])

# 15　打印输出最终的测试结果。（7分）
print('Testing')
model.evaluate(ds_test)
print('Tested. Over.')
